import os
from nlp_tools.corpus.seq2seq.autotitle_thunc import AutoTitleThunc
from nlp_tools.tasks.text_generator.bert_generator_model import BertGeneratorModel
from nlp_tools.processors.sequence_processor import SequenceProcessor
from nlp_tools.tokenizer.bert_tokenizer import BertTokenizer
from nlp_tools.utils.text_generator_utils import AutoTitle

from tensorflow.keras.callbacks import Callback

model_save_path = '/home/qiufengfeng/nlp/train_models/text_autogeneter/'
thunc_data_path = '/home/qiufengfeng/nlp/nlp_data/classify/THUCNews'
#x,y = AutoTitleThunc.load_data(thunc_data_path)


bert_model_path = "/home/qiufengfeng/nlp/pre_trained_model/chinese_L-12_H-768_A-12/chinese_L-12_H-768_A-12"

from bert4keras.tokenizers import load_vocab
token_dict, keep_tokens = load_vocab(
    dict_path=os.path.join(bert_model_path,'vocab.txt'),
    simplified=True,
    startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]'],
)

text_tokenizer  = BertTokenizer(token_dict=os.path.join(bert_model_path,'vocab.txt'),simplified=True)
text_processor = SequenceProcessor(text_tokenizer=text_tokenizer)
model = BertGeneratorModel(bert_model_path=bert_model_path,text_processor=text_processor,sequence_length=256)




autotitle = AutoTitle(model,text_tokenizer,start_id=None, end_id=text_tokenizer._token_end_id,max_generator_len=32,max_seq_len=256)

def just_show():
    s1 = u'夏天来临，皮肤在强烈紫外线的照射下，晒伤不可避免，因此，晒后及时修复显得尤为重要，否则可能会造成长期伤害。专家表示，选择晒后护肤品要慎重，芦荟凝胶是最安全，有效的一种选择，晒伤严重者，还请及 时 就医 。'
    s2 = u'8月28日，网络爆料称，华住集团旗下连锁酒店用户数据疑似发生泄露。从卖家发布的内容看，数据包含华住旗下汉庭、禧玥、桔子、宜必思等10余个品牌酒店的住客信息。泄露的信息包括华住官网注册资料、酒店入住登记的身份信息及酒店开房记录，住客姓名、手机号、邮箱、身份证号、登录账号密码等。卖家对这个约5亿条数据打包出售。第三方安全平台威胁猎人对信息出售者提供的三万条数据进行验证，认为数据真实性非常高。当天下午 ，华 住集 团发声明称，已在内部迅速开展核查，并第一时间报警。当晚，上海警方消息称，接到华住集团报案，警方已经介入调查。'
    for s in [s1, s2]:
        print(u'生成标题:', autotitle.generate(s))

class SaveAndTextEvaluator(Callback):
    """评估与保存
    """
    def __init__(self,**kwargs):
        super(SaveAndTextEvaluator,self).__init__()
        self.lowest = 1e10
        self.model = model

    def on_epoch_end(self, epoch, logs=None):
        # 保存最优
        if logs['loss'] <= self.lowest:
            self.lowest = logs['loss']
            model.save(model_save_path)
        # 演示效果
        just_show()
thunc_data = AutoTitleThunc(file_path=thunc_data_path)
model.fit_generator(thunc_data,epochs=20,batch_size=16,callbacks=[SaveAndTextEvaluator()],fit_kwargs={"steps_per_epoch":200})
#x,y = thunc_data.load_data()
#model.fit(x,y,epochs=20,batch_size=16,callbacks=[SaveAndTextEvaluator()],fit_kwargs={"steps_per_epoch":2})
